package main

import (
	"github.com/sirupsen/logrus"
	"os"
	"os/signal"
	"simPSapp/pkg/server"
	"simPSapp/pkg/worker"
	"strconv"
)

func main() {
	mode := os.Getenv("MODE")
	// size的单位为MB
	dataSize, err := strconv.Atoi(os.Getenv("DATA_SIZE"))
	if err != nil {
		logrus.Fatalf("[main]-无法解析DATA_SIZE: %v", err)
	}
	workerNum, err := strconv.Atoi(os.Getenv("WORKER_NUM"))
	if err != nil {
		logrus.Fatalf("[main]-无法解析WORKER_NUM: %v", err)
	}
	if mode == "" {
		mode = "server"
	}

	port := os.Getenv("PORT")

	if port == "" {
		port = "8080"
	}

	if mode == "server" {
		logrus.Infof("[main]-start as %s", mode)
		sv := server.NewServer(dataSize, workerNum, port)
		go sv.Run()
		server.WaitForGrpc(sv)
		stopCh := make(chan os.Signal, 1)
		signal.Notify(stopCh, os.Interrupt, os.Kill)
		<-stopCh
		sv.Stop()

	} else if mode == "worker" {
		logrus.Infof("[main]-start as %s", mode)
		serverAddr := os.Getenv("SERVER_ADDR")
		batchSize, err := strconv.Atoi(os.Getenv("BATCH_SIZE"))
		if err != nil {
			logrus.Fatalf("[main]-无法解析BATCH_SIZE: %v", err)
		}
		gradientSize, err := strconv.Atoi(os.Getenv("GRADIENT_SIZE"))
		if err != nil {
			logrus.Fatalf("[main]-无法解析GRADIENT_SIZE: %v", err)
		}
		workerId, err := strconv.Atoi(os.Getenv("WORKER_ID"))
		if err != nil {
			logrus.Fatalf("[main]-无法解析WORKER_ID: %v", err)
		}
		worker.NewWorker(serverAddr, dataSize, batchSize, gradientSize, workerId).Run()
	} else {
		logrus.Fatalf("[main]-mode = %s错误，应该标识为server或者worker", mode)
	}
}
